# This code is part of cqlib.
#
# (C) Copyright China Telecom Quantum Group, QuantumCTek Co., Ltd.,
# Center for Excellence in Quantum Information and Quantum Physics 2024.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
GuoDun quantum platform tests
"""
# pylint: disable=redefined-outer-name

import os
import time
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import functools
import pytest
import random
from cqlib import GuoDunPlatform
from cqlib.quantum_platform.base import QuantumLanguage
from cqlib.exceptions import CqlibRequestError

login_key = os.getenv('LOGIN_KEY')
sim_machine = os.getenv('SIM_MACHINE')
quantum_machine = os.getenv('QUANTUM_MACHINE')


@pytest.fixture
def platform():
    """
    Fixture to create an instance of the GuoDunPlatform class
    """
    if scheme := os.getenv('SCHEME'):
        GuoDunPlatform.SCHEME = scheme
    if domain := os.getenv('DOMAIN'):
        GuoDunPlatform.DOMAIN = domain
    assert login_key, 'Please set the login key first'
    assert sim_machine, 'Please set the simulator code first'
    assert quantum_machine, 'Please set the quantum computer code first'

    return GuoDunPlatform(login_key=login_key, machine_name=quantum_machine)


@pytest.fixture
def lab_id(platform):  # pylint: disable=redefined-outer-name
    """Create a lab collection for test cases"""
    name = f'lab.{datetime.now().strftime("%Y%m%d%H%M%S.%f")}.{random.randint(0, 100000)}'
    return platform.create_lab(name)


class TestGuoDunPlatformBasic:
    """
    Test cases for the GuoDunPlatform
    """

    def test_login(self, platform):
        """
        test platform login
        """
        print(f'access_token: {platform.access_token}')
        assert platform.access_token

    def test_query_quantum_computer_list(self, platform):
        """
        test platform query quantum computer
        """
        computer_list = platform.query_quantum_computer_list()
        print(f'computer_list: {computer_list}')
        assert isinstance(computer_list, list)

    def test_send_request_with_expired_token(self, platform):
        """
        After login, the access token will expire after 24 hours.
        Test if platform will re-login given expired token.
        """
        expired_token = 'test_expired_token'
        lab_name = f'lab.{datetime.now().strftime("%Y%m%d%H%M%S.%f")}.expired_token'
        platform.access_token = expired_token
        new_lab_id = platform.create_lab(lab_name, remark='test_expired_token')
        assert new_lab_id

    def test_download_config(self, platform):
        """
        test platform download quantum computer config
        """
        cfg = platform.download_config()
        print(f'cfg: {cfg}')
        assert cfg

    def test_qcis_check_regular_false(self, platform):
        """
        check if given qcis order with wrong grammar, false will return
        """
        platform.set_machine(quantum_machine)
        data = platform.qcis_check_regular('HQ0\nH Q2\nCZ Q0 Q2\nM Q0\nM Q2')
        print(f'data: {data}')
        assert data is False

    def test_qcis_check_regular_quantum_machine_false(self, platform):
        """
        check if given invalid mapping qcis order, false will return
        """
        platform.set_machine(quantum_machine)
        data = platform.qcis_check_regular('H Q0\nH Q2\nCZ Q0 Q2\nM Q0\nM Q2')
        print(f'data: {data}')
        assert data is False

    def test_qcis_check_regular_quantum_machine_true(self, platform):
        """
        check if given valid mapping qcis order, true will return
        """
        platform.set_machine(quantum_machine)
        data = platform.qcis_check_regular('H Q0\nH Q7\nCZ Q0 Q7\nM Q0\nM Q7')
        platform.set_machine(sim_machine)
        print(f'data: {data}')
        assert data


class TestGuoDunPlatformQcis:
    """
    test guodun platform with qcis circuit
    """

    qcis_circuit = """
x q1
m q3
"""

    def test_save_experiment_qcis(self, platform, lab_id):
        """
        test platform save qcis experiment given lab id
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}')
        print(f'exp_id: {exp_id}')
        assert exp_id

    def test_run_experiment_qcis(self, platform, lab_id):
        """
        test platform run qcis experiment given lab id
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job(self, platform, lab_id):
        """
        test platform submit job given lab id
        """
        query_id = platform.submit_job(
            circuit=self.qcis_circuit,
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            num_shots=50,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_multi_thread(self, platform, lab_id, count=10):
        """
        test platform submit job given lab id and send requests with multi thread
        """
        partial_submit_job = functools.partial(platform.submit_job,
                                               circuit=self.qcis_circuit,
                                               lab_id=lab_id,
                                               num_shots=50,
                                               )

        executor = ThreadPoolExecutor(max_workers=count)
        all_tasks = [executor.submit(partial_submit_job, exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}_{i}')
                     for
                     i in range(count)]
        for future in as_completed(all_tasks):
            data = future.result()
            assert data

    def test_submit_job_batch(self, platform, lab_id, circuit_num=30):
        """
        test platform submit job given lab id and send multiple qcis in one batch
        """
        circuits_list = [self.qcis_circuit for _ in range(circuit_num)]
        query_id = platform.submit_job(
            circuit=circuits_list,
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            num_shots=50,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_exp_circuit(self, platform, lab_id):
        """
        test platform submit job given exp id
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        # given qcis code and exp_id at the same time, will throw error
        with pytest.raises(CqlibRequestError):
            query_id = platform.submit_job(
                circuit=self.qcis_circuit,
                exp_id=exp_id,
                num_shots=50,
            )

    def test_submit_job_lab_exp(self, platform, lab_id):
        """
        test platform submit job given both lab id and exp id
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.submit_job(
            exp_id=exp_id,
            lab_id=lab_id,
            num_shots=50,
            language=QuantumLanguage.QCIS,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_parameters(self, platform, lab_id):
        """
        test platform submit job given lab id with parameters
        """
        circuit = "RX Q0 {theta}\nM Q0"
        query_id = platform.submit_job(
            circuit=circuit,
            parameters=[['theta']],
            values=[[1.2]],
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            num_shots=50,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_exp_id(self, platform, lab_id):
        """
        test platform submit job given exp id
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        query_id = platform.submit_job(
            exp_id=exp_id,
            num_shots=50,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_query_experiment(self, platform, lab_id):
        """
        test platform submit experiment and query experiment
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        exp = platform.query_experiment(query_id)
        print(f'exp: {exp}')
        assert exp

    def test_get_experiment_circuit(self, platform, lab_id):
        """
        test platform submit experiment and get experiment circuit
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        data = platform.get_experiment_circuit(query_id)
        print(f'circuit: {data}')
        assert data

    def test_submit_experiment_lab_id(self, platform, lab_id):
        """
        test platform submit experiment given lab id using circuits with parameters
        """
        circuit = "RX Q0 {theta}\nM Q0"
        query_id = platform.submit_job(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            parameters=[['theta']],
            values=[[1.2]],
            lab_id=lab_id,
            num_shots=50,
        )
        assert query_id

    def test_submit_experiment_parameters(self, platform):
        """
        test platform submit experiment without lab id using circuits with parameters
        """
        circuit = "RX Q0 {theta}\nM Q0"
        query_id = platform.submit_job(
            circuit=circuit,
            language=QuantumLanguage.QCIS,
            parameters=[['theta']],
            values=[[1.2]],
            num_shots=50,
        )
        assert query_id


class TestGuoDunPlatformIsq:
    """
    test guodun platform with isq circuit
    """
    isq_circuit = """
//实现两个qubit之间的交换
qbit q[66];
procedure main() {

X<q[0]>;

H<q[6]>;
CZ<q[0], q[6]>;
H<q[6]>;
H<q[0]>;
CZ<q[6], q[0]>;
H<q[0]>;
H<q[6]>;
CZ<q[0], q[6]>;
H<q[6]>;

M<q[6]>;

}
"""

    def test_save_experiment_isq(self, platform, lab_id):
        """
        test platform save experiment given isq circuit
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        assert exp_id

    def test_run_experiment_isq(self, platform, lab_id):
        """
        test platform save and run experiment given isq circuit
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_isq_exp_circuit(self, platform, lab_id):
        """
        test platform save experiment and submit job given isq circuit
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.submit_job(
            exp_id=exp_id,
            language=QuantumLanguage.ISQ,
            num_shots=50,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_parameters_isq(self, platform, lab_id):
        """
        test platform submit job with lab id given isq circuit and parameters
        """
        query_id = platform.submit_job(
            circuit=[self.isq_circuit],
            exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
            lab_id=lab_id,
            language=QuantumLanguage.ISQ,
            num_shots=50,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_submit_job_isq_2(self, platform):
        """
        test platform submit job without lab id given isq circuit
        """
        query_id = platform.submit_job(
            circuit=[self.isq_circuit],
            num_shots=50,
            language=QuantumLanguage.ISQ,
        )
        print(f'query_id: {query_id}')
        assert query_id

    def test_query_experiment_isq(self, platform, lab_id):
        """
        test platform submit experiment and query experiment given isq circuit
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        exp = platform.query_experiment(query_id)
        print(f'exp: {exp}')
        assert exp

    def test_get_experiment_circuit_isq(self, platform, lab_id):
        """
        test platform submit experiment and get experiment circuit given isq circuit
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.isq_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.ISQ)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        data = platform.get_experiment_circuit(query_id)
        print(f'circuit: {data}')
        assert data


class TestGuoDunPlatformReExecuteTask:
    """
    test guodun platform re execute task
    """
    qcis_circuit = """
x q1
m q3
"""

    def test_re_execute_task_lab_id_input(self, platform, lab_id):
        """
        test platform re-execute task given lab id only
        """
        print(f'lab_id: {lab_id}')
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        # time.sleep(5)
        data = platform.re_execute_task(lab_id=lab_id, query_id=None)
        print(f'data: {data}')
        assert data

    def test_re_execute_task_query_id_input(self, platform, lab_id):
        """
        test platform re-execute task given query id only
        """
        print(f'lab_id: {lab_id}')
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=50)
        print(f'query_id: {query_id}')
        # time.sleep(5)
        data = platform.re_execute_task(lab_id=None, query_id=query_id[0])
        print(f'data: {data}')
        assert data

    def test_re_execute_task_large_shots(self, platform, lab_id):
        """
        test platform re-execute task given large shots
        """
        print(f'lab_id: {lab_id}')
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=10000)
        print(f'query_id: {query_id}')
        data = platform.re_execute_task(lab_id=None, query_id=query_id[0])
        print(f'data: {data}')
        assert data


class TestGuoDunPlatformStopRunningExperiments:
    """
    test guodun platform stop running experiments
    """
    qcis_circuit = """
x q1
m q3
"""

    def test_stop_running_experiments_query_id_input(self, platform, lab_id):
        """
        test platform stop running experiments given query id only
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=100)
        print(f'query_id: {query_id}')
        # data = platform.stop_running_experiments(lab_id=lab_id, query_id=None)
        data = platform.stop_running_experiments(lab_id=None, query_id=query_id[0])
        print(f'data: {data}')
        assert data

    def test_stop_running_experiments_lab_id_input(self, platform, lab_id):
        """
        test platform stop running experiments given lab id only
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=100)
        print(f'query_id: {query_id}')
        data = platform.stop_running_experiments(lab_id=lab_id, query_id=None)
        print(f'data: {data}')
        assert data

    def test_stop_running_experiments_small_shots(self, platform, lab_id):
        """
        test platform stop running experiments given small shots
        """
        exp_id = platform.save_experiment(lab_id=lab_id, circuit=self.qcis_circuit,
                                          name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
                                          language=QuantumLanguage.QCIS)
        print(f'exp_id: {exp_id}')
        query_id = platform.run_experiment(exp_id=exp_id, num_shots=20)
        print(f'query_id: {query_id}')
        # data = platform.stop_running_experiments(lab_id=lab_id, query_id=None)
        data = platform.stop_running_experiments(lab_id=None, query_id=query_id[0])
        print(f'data: {data}')
        assert data
